'''
@Time    : 2022/3/22 20:26
@Author  : Fu Junyu
@Site    : www.fujunyu.cn
@File    : HeteroRGCNLayer.py
@Software: PyCharm
'''
import torch
import torch.nn as nn
import dgl.function as fn

'''
构建一个RGCN层
'''
class HeteroRGCNLayer(nn.Module):
    def __init__(self, user_in_size, movies_in_size, out_size, num_bases=3, num_rels=3, regularizer='basis'):
        super(HeteroRGCNLayer, self).__init__()

        self.user_in_size = user_in_size
        self.movies_in_size = movies_in_size
        self.out_size =out_size
        self.num_bases = num_bases
        self.num_rels =num_rels
        self.regularizer = regularizer
        self.etypeDic = {}
        self.user_weight = None
        self.movies_weight = None

        if self.regularizer == "basis":
            # add basis weights
            self.user_weight = nn.Parameter(torch.Tensor(self.num_bases, self.user_in_size, self.out_size))
            self.movies_weight = nn.Parameter(torch.Tensor(self.num_bases, self.movies_in_size, self.out_size))
            if self.num_bases < self.num_rels:
                # linear combination coefficients
                self.w_comp = nn.Parameter(torch.Tensor(self.num_rels, self.num_bases))
            nn.init.xavier_uniform_(self.user_weight, gain=nn.init.calculate_gain('sigmoid'))
            nn.init.xavier_uniform_(self.movies_weight, gain=nn.init.calculate_gain('sigmoid'))
            if self.num_bases < self.num_rels:
                nn.init.xavier_uniform_(self.w_comp, gain=nn.init.calculate_gain('sigmoid'))

        self.norm_movies = nn.BatchNorm1d(self.out_size)
        self.Norm_user = nn.BatchNorm1d(self.out_size)

    def forward(self, graph, feat_dict):
        user_weight, movies_weight = self.basis()

        for i, etype in enumerate(graph.etypes):
            self.etypeDic[etype] = i

        funcs = {}

        act = nn.Sigmoid()

        for srctype, etype, dsttype in graph.canonical_etypes:
            if dsttype == 'user':
                Wh = self.norm_movies(act(torch.mm(feat_dict[srctype], movies_weight[self.etypeDic[etype]])))
                graph.nodes[srctype].data['Wh_%s' % etype] = Wh
                funcs[etype] = (fn.copy_u('Wh_%s' % etype, 'm'), fn.mean('m', 'h'))

            elif dsttype == 'movies':
                Wh = self.Norm_user(act(torch.mm(feat_dict[srctype], user_weight[self.etypeDic[etype]])))
                graph.nodes[srctype].data['Wh_%s' % etype] = Wh
                funcs[etype] = (fn.copy_u('Wh_%s' % etype, 'm'), fn.mean('m', 'h'))

        graph.multi_update_all(funcs, 'sum')

        return {ntype: graph.nodes[ntype].data['h'] for ntype in graph.ntypes}

    def basis(self):
        user_weight = None
        movies_weight = None

        if self.num_bases < self.num_rels:
            # generate all weights from bases
            user_weight = self.user_weight.view(self.num_bases, self.user_in_size * self.out_size)
            user_weight = torch.matmul(self.w_comp, user_weight).view(self.num_rels, self.user_in_size, self.out_size)

            movies_weight = self.movies_weight.view(self.num_bases, self.movies_in_size * self.out_size)
            movies_weight = torch.matmul(self.w_comp, movies_weight).view(self.num_rels, self.movies_in_size, self.out_size)
        else:
            user_weight = self.user_weight
            movies_weight = self.movies_weight

        return (user_weight, movies_weight)







